"""
This files implements the algorithm to calculate a gradient that evens out
a given retention time distribution
"""


import os
import sys 
import math
import matplotlib.pyplot as plt
import argparse


class Peptide:
    """
    Class to store information about a peptide
    """
    
    def __init__(self, seq, rt, perc_b):
        self.seq = seq
        self.rt = float(rt)
        self.perc_b = float(perc_b)
        

    def __str__(self):
        return "%s:%f:%f" % (self.seq, self.rt, self.perc_b)
        

def get_perc_b(rt, start_gradient, end_gradient):
    """
    Input: 
    rt - a retention time expressed as raw file time
    start_gradient - (time, %B)  when the gradient starts; the time is 
            raw file time
    end_gradient - (time, %B) when the gradient ends; the time is raw 
            file time
    
    Output: assuming a linear gradient between start_gradient and 
            end_gradient, return the %B corresponding to the time rt
    """
    
    rt_interval = end_gradient[0]-start_gradient[0]
    return (((rt-start_gradient[0]) * \
            (end_gradient[1]-start_gradient[1])) / rt_interval) + \
            start_gradient[1]


def load_rt_file(rt_file, start_gradient, end_gradient):
    """
    Input:
    rt_file - file including an id, retention time, one such pair per 
            line, space-separated. The id can be a peptide sequence, 
            or any other text
    start_gradient - (time, %B)  when the gradient starts; the time is 
            raw file time
    end_gradient - (time, %B) when the gradient ends; the time is raw 
            file time
            
    Output: a list of Peptide objects loaded from the rt_file. Only those
            peptides that elute between start and end gradient are 
            included
    """
    
    print "\nLoading {0} ...".format(rt_file)
    lines = filter(lambda l: not l.startswith("#"), \
            open(rt_file, "r").readlines())
    peptides = []
    for l in lines:
        fields = l.split()
        rt = float(fields[1])
        if rt >= start_gradient[0] and rt <= end_gradient[0]:
            b = get_perc_b(rt, start_gradient, end_gradient)
            peptides.append(Peptide(fields[0], rt, b))
    print "{0} peptides, {1} elute during the gradient time".format( \
            len(lines), len(peptides))
  
    return peptides
    

def load_linear_gradient(linear_gradient_file, lag_time):
    """
    Input: 
    linear_gradient_file - a file including a linear gradient in LC times
            This files should include the following two lines:
            start_gradient = 16.2, 2.0
            end_gradient = 256.2, 40.0
    lag_time- the value that should be added to the LC time to get the
            raw file time 
            
    Return: the start and end of the gradient, each given as a pair 
            (retention time, %B)
    """
    
    print "\nLoading linear gradient from {0}".format(\
            linear_gradient_file)
    lines = [l for l in open(linear_gradient_file).readlines() \
        if not l.startswith("#") ]
 
    start_gradient = (None, None)
    end_gradient = (None, None)
    s, e = False, False
    for l in lines:
        if l.startswith("start_gradient"):
            tmp = l.split("=")[1].strip()
            start_gradient = (float(tmp.split(",")[0]) + lag_time, \
                    float(tmp.split(",")[1]))
            s = True
        if l.startswith("end_gradient"):
            tmp = l.split("=")[1].strip()
            end_gradient = (float(tmp.split(",")[0]) + lag_time, \
                    float(tmp.split(",")[1]))
            e = True
 
    if not s or not e:
        sys.exit("Incorrect gradient file. Check the format of the " + \
                "gradient file and try again")
    print "Raw file times, gradient start={0}, gradient end={1}".format(\
            start_gradient, end_gradient)
    return start_gradient, end_gradient


def calculate_nonlinear_gradient(peptides, start_gradient, \
        end_gradient, step_size):
    """
    Input:
    peptides - list of Peptide objects whose distribution will be evened 
            by the nonlinear gradient 
    start_gradient - (time, %B)  when the gradient starts; the time is 
            raw file time
    end_gradient - (time, %B) when the gradient ends; the time is raw 
            file time
    step_size - how often should we sample the gradient 
    
    Output: the nonlinear gradient that evens out the list of peptides, 
            expressed as a list of pairs (retention time, %B) 
    """
    print "\nCalculate nonlinear gradient ..."
    # compute how many peptides are per time unit
    n_peptides = float(len(peptides)) / \
            ((end_gradient[0]-start_gradient[0])/step_size)
    print "%d peptides per time interval" % (n_peptides, )  
    peptides.sort(key = lambda pep: pep.rt)    
    
    # compute gradient 
    gradient = [start_gradient]
    i = n_peptides
    rt = start_gradient[0]+step_size
    j = 1
    while i < len(peptides)-1: 
        if isinstance(i, int): 
            #gradient.append((rt, peptides[i].perc_b))
            b = peptides[i].perc_b
        else:
            prev_index = int(math.floor(i))
            next_index = int(math.ceil(i))
            prev_b = peptides[prev_index].perc_b
            next_b = peptides[next_index].perc_b
            b = (i-prev_index)*prev_b + (next_index-i)*next_b
            # this could happen because of rounding errors 
            if b < prev_b:
                b = prev_b                                            
        gradient.append((rt, b))
        rt += step_size
        i += n_peptides
    
    gradient.append(end_gradient)    
    print "Gradient calculated."        
    return gradient 


def plot_gradient(gradient, start_gradient, end_gradient, out_file):
    """
    Input:
    gradient - the nonlinear gradient given as a list of pairs 
            (retention time, %B)
    start_gradient - (time, %B)  when the linear gradient starts; 
            the time is raw file time
    end_gradient - (time, %B) when the linear gradient ends; 
            the time is raw file time
    out_file - file to plot the gradient to (no extension)
    
    Output: the gradients are plotted to out_file
    """
    
    rts = [p[0] for p in gradient]
    plt.plot(rts, [get_perc_b(rt, start_gradient, end_gradient) \
            for rt in rts], '-', label = "Linear gradient")  
    plt.plot(rts, [p[1] for p in gradient], '--', \
            label = "Nonlinear gradient")  
    plt.xlabel("Raw file retention time [min]")
    plt.ylabel("%B")
    plt.legend(loc = "upper left")
    plt.savefig(out_file, format = 'png')
    plt.cla()
    plt.close()
    

def write_gradient_file(gradient, out_file, decimals, remove_redundancy):
    """
    Input:
    gradient - the nonlinear gradient given as a list of pairs 
            (retention time, %B)
    out_file - name of the output file 
    decimals - the number of decimals used to give the %B
    remove_redundancy - do not includ multiple lines with the same %B
    
    Output: the gradient will be written to the output file 
    """
    print "Decimals:", decimals
    print "\nWriting gradient to {0}".format(out_file)
    outf = open(out_file, "w")
    outf.write("Time(LC-time)\t%B\n")
    prev = None
    for p in gradient:
        b = round(p[1], decimals)
        if (not remove_redundancy) or (prev == None) or (prev != b):
            outf.write('{0}\t{1}\n'.format(p[0], b))
        prev = b
    outf.close()
    print "Done."


def print_command_line(args):
    print "\n############################################"    
    names = ["linear_gradient_file", "rt_file", "lag_time", \
            "out_file", "step_size", "n_decimals", "remove_redundancy"]
    arguments = ["g", "r", "l", "o", "s", "d", "c"]
    for name, arg in zip(names, arguments):
        data = str(eval("args." + arg))
        print "{0} = {1}".format(name, data)    
    print "############################################"  
    
    
def string_to_bool(s):
    if s.lower() == "t" or s.lower() == "true":
        return True
    return False
    
    
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-g', help='Linear gradient file. Needs '+\
            'to include two lines like: "start_gradient = 10.0, 2.0", and '+\
            '"end_gradient = 250.0, 40.0", where 10.0 and 250.0 should be '+\
            'replaced by the start and end times of the gradient as LC times,'+\
            ' and 2.0 and 40.0 should be replaced by the start and end '+\
            'of the gradient in percentage B',\
            metavar = "linear_gradient_file",\
            required = True)
    parser.add_argument('-r', help='File including the distribution '+\
            'that the nonlinear gradient aims at uniformizing. Each '+\
            'line of this file should include an id, and a retention '+\
            'time, space or tab separated. These are expected to be '+\
            'expressed in raw file time', metavar = "rt_file", \
            required = True)            
    parser.add_argument('-l', help='Value to be rested from the raw ' +\
            'file time to get the LC time. Note that the gradients '+\
            'are expressed in LC times, while the peptide retention '+\
            'times are given in raw file times', metavar = "lag_time", \
            required = True)
    parser.add_argument('-o', help='Output file (no extension). '+\
            'Default: nonlinear_gradient', metavar = "out_file", \
            required = False, default="nonlinear_gradient")
    parser.add_argument('-s', help='The step size when calculating '+\
            'the nonlinear gradient. Default=1.0', \
            metavar = "step_size", required = False, default = 1.0)
    parser.add_argument('-d', help='The number of decimals used to '+\
            'output the nonlinear gradient. Default=3', \
            metavar = "number_decimals", required = False, default = 3)    
    parser.add_argument('-c', help='Remove redundancy in the gradient'+\
            '[T/F]. Default = T', metavar = "remove_redundancy", \
            required = False, default = "T")                
            

    args = parser.parse_args()
    linear_gradient_file = args.g
    rt_file = args.r
    lag_time = float(args.l)
    out_file = args.o
    step_size = float(args.s)
    decimals = int(args.d)  
    remove_redundancy = string_to_bool(args.c)
    print_command_line(args) 
  
    # load data  
    start_gradient, end_gradient = load_linear_gradient(\
            linear_gradient_file, lag_time)
    peptides = load_rt_file(rt_file, start_gradient, end_gradient)
    
    # design the nonlinear gradient
    nonlinear_gradient = calculate_nonlinear_gradient(peptides, \
            start_gradient, end_gradient, step_size)
   
    # plot and write the gradient to the output file 
    plot_gradient(nonlinear_gradient, start_gradient, end_gradient, \
            out_file + ".png")
    nonlinear_lc_gradient = [(p[0]-lag_time, p[1]) \
            for p in nonlinear_gradient]           
    write_gradient_file(nonlinear_lc_gradient, out_file + ".txt", \
            decimals, remove_redundancy)
    

if __name__ == '__main__':
    main()
    
